/*
 * Copyright (c) 2023 Huawei Device Co., Ltd.
 * Licensed under the Apache License, Version 2.0 (the "License");
 * you may not use this file except in compliance with the License.
 * You may obtain a copy of the License at
 *
 *     http://www.apache.org/licenses/LICENSE-2.0
 *
 * Unless required by applicable law or agreed to in writing, software
 * distributed under the License is distributed on an "AS IS" BASIS,
 * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
 * See the License for the specific language governing permissions and
 * limitations under the License.
 */



#include <linux/bpf.h>
#include <linux/pkt_cls.h>
#include "bpf/bpf_helpers.h"
#include "bpf/bpf_endian.h"
#include "bpf_def.h"

#include <linux/if_packet.h>
#include <linux/if_ether.h>
#include <linux/ip.h>
#include <linux/tcp.h>
#include <linux/in.h>


#ifndef offsetof
#define offsetof(TYPE, MEMEBER) ((size_t) &((TYPE *)0)->MEMBER)
#endif

#define IS_PSEUDO 0x10

#define m_iSeqNoTH 0x7FFFFFFF
#define m_iMaxSeqNo 0xFFFFFFFF


bpf_map_def SEC("maps") conn_loss_stats_map = {
    .type = BPF_MAP_TYPE_HASH,
    .key_size = sizeof(conn_loss_stat_key),
    .value_size = sizeof(conn_loss_stat_value),
    .max_entries = 1024,
    .map_flags = 0,
    .inner_map_idx = 0,
    .numa_node = 0,
};

bpf_map_def SEC("maps") conn_rcv_max_seq = {
    .type = BPF_MAP_TYPE_HASH,
    .key_size = sizeof(conn_loss_stat_key),
    .value_size = sizeof(__u32),
    .max_entries = 1,
    .map_flags = 0,
    .inner_map_idx = 0,
    .numa_node = 0,
};

inline static unsigned int cal_abs(long long int num) {
    return num >= 0? num:-num;
}

inline static long long int cal_seqoff(__u32 seq1, __u32 seq2)
{
    if (cal_abs(seq1 - seq2) < m_iSeqNoTH) {
        return seq2 - seq1;
    }

    if (seq1 < seq2) {
        return seq2 - seq1 - m_iMaxSeqNo - 1;
    }

    return seq2 - seq1 + m_iMaxSeqNo + 1;
}


SEC("classifier/egress_window_update")
int egress_window_update(struct __sk_buff *skb)
{
    int ret = 0;
    struct ethhdr l2;
    struct iphdr l3;
    struct tcphdr l4;

    if (skb->protocol != bpf_htons(ETH_P_IP)) {
        return TC_ACT_OK;
    }
    
    // Load L2 (Ethernet) header
    ret = bpf_skb_load_bytes(skb, 0, &l2, sizeof(l2));
    if (ret < 0) {
        return TC_ACT_OK;
    }

    // Load L3 (IP) header
    ret = bpf_skb_load_bytes(skb, sizeof(l2), &l3, sizeof(l3));
    if (ret < 0) {
        return TC_ACT_OK;
    }

    // Check if it's TCP
    if (l3.protocol != IPPROTO_TCP) {
        return TC_ACT_OK;
    }

    // Load L4 (TCP) header
    ret = bpf_skb_load_bytes(skb, sizeof(l2) + l3.ihl * 4, &l4, sizeof(l4));
    if (ret < 0) {
        return TC_ACT_OK;
    }

    __u32 tcp_offset = sizeof(l2) + l3.ihl * 4;
    __u16 old_window = 0;
    __u16 old_check = 0;
    __u16 new_window = bpf_htons(1024);
    __u16 new_check = 0;
    if (bpf_ntohs(l4.dest) == 9999) {
        ret = bpf_skb_load_bytes(skb, tcp_offset + offsetof(struct tcphdr, window), &old_window, sizeof(old_window));
        if (ret < 0) {
            bpf_printk("bpf_skb_load_bytes old window failed ret = %d", ret);
            return TC_ACT_OK;
        }
        ret = bpf_skb_load_bytes(skb, tcp_offset + offsetof(struct tcphdr, check), &old_check, sizeof(old_check));
        if (ret < 0) {
            bpf_printk("bpf_skb_load_bytes old check failed ret = %d", ret);
            return TC_ACT_OK;
        }
        bpf_printk("old window = %x, check = %x", old_window, old_check);
        ret = bpf_l4_csum_replace(skb, tcp_offset + offsetof(struct tcphdr, check), old_window, new_window, IS_PSEUDO | sizeof(new_window));
		if (ret < 0) {
			bpf_printk("bpf_l4_csum_replace window failed: %d\n");
			return TC_ACT_OK;
		}
        // 修改窗口
        ret = bpf_skb_store_bytes(skb, tcp_offset + offsetof(struct tcphdr, window), &new_window, sizeof(new_window), 0);
        if (ret < 0) {
            bpf_printk("bpf_skb_store_bytes window failed ret = %d", ret);
            return TC_ACT_OK;
        }
        ret = bpf_skb_load_bytes(skb, tcp_offset + offsetof(struct tcphdr, check), &new_check, sizeof(new_check));
        if (ret < 0) {
            bpf_printk("bpf_skb_load_bytes new check failed ret = %d", ret);
        }
        bpf_printk("new window = %x, check = %x", new_window, new_check);
    }

    return TC_ACT_OK;
}


SEC("classifier/ingress_pkt_redirect")
int ingress_pkt_redirect(struct __sk_buff *skb)
{
    int ret = 0;
    struct ethhdr l2;
    struct iphdr l3;
    struct tcphdr l4;

    if (skb->protocol != bpf_htons(ETH_P_IP)) {
        return TC_ACT_OK;
    }
    
    // Load L2 (Ethernet) header
    ret = bpf_skb_load_bytes(skb, 0, &l2, sizeof(l2));
    if (ret < 0) {
        return TC_ACT_OK;
    }

    // Load L3 (IP) header
    ret = bpf_skb_load_bytes(skb, sizeof(l2), &l3, sizeof(l3));
    if (ret < 0) {
        return TC_ACT_OK;
    }

    // Check if it's TCP
    if (l3.protocol != IPPROTO_TCP) {
        return TC_ACT_OK;
    }

    // Load L4 (TCP) header
    ret = bpf_skb_load_bytes(skb, sizeof(l2) + l3.ihl * 4, &l4, sizeof(l4));
    if (ret < 0) {
        return TC_ACT_OK;
    }

    if (bpf_ntohs(l4.dest) != 9999) {
        return TC_ACT_OK;
    }

    __u32 tcp_offset = sizeof(l2) + l3.ihl * 4;

    if (l4.syn == 1 && l4.res1 == 0) {
        bpf_printk("1111111111");
        l4.res1 = 1;
        bpf_skb_store_bytes(skb, tcp_offset, &l4, sizeof(l4), 0);
        bpf_clone_redirect(skb, skb->ifindex, BPF_F_INGRESS);
    }

    return TC_ACT_OK;
}


SEC("classifier/ingress_losspkt_statistic")
int ingress_losspkt_statistic(struct __sk_buff *skb)
{
    int ret = 0;
    struct ethhdr l2;
    struct iphdr l3;
    struct tcphdr l4;

    if (skb->protocol != bpf_htons(ETH_P_IP)) {
        return TC_ACT_OK;
    }
    
    // Load L2 (Ethernet) header
    ret = bpf_skb_load_bytes(skb, 0, &l2, sizeof(l2));
    if (ret < 0) {
        return TC_ACT_OK;
    }

    // Load L3 (IP) header
    ret = bpf_skb_load_bytes(skb, sizeof(l2), &l3, sizeof(l3));
    if (ret < 0) {
        return TC_ACT_OK;
    }

    // Check if it's TCP
    if (l3.protocol != IPPROTO_TCP) {
        return TC_ACT_OK;
    }

    // Load L4 (TCP) header
    ret = bpf_skb_load_bytes(skb, sizeof(l2) + l3.ihl * 4, &l4, sizeof(l4));
    if (ret < 0) {
        return TC_ACT_OK;
    }

    if (bpf_ntohs(l4.dest) != 9999) {
        return TC_ACT_OK;
    }

    __u32 tcplen = bpf_ntohs(l3.tot_len) - l3.ihl * 4;
    __u32 tcphdrlen = l4.doff * 4;
    __u32 tcp_data_len = tcplen - tcphdrlen;

    if (tcp_data_len <= 0) {
        return TC_ACT_OK;
    }

    conn_loss_stat_key clskey = {
        .saddr = l3.saddr,
        .daddr = l3.daddr,
        .sport = l4.source,
        .dport = l4.dest,
        .protocol = l3.protocol
    };

    conn_loss_stat_value *clsvalue = bpf_map_lookup_elem(&conn_loss_stats_map, &clskey);

    bpf_printk("*********");
    bpf_printk("tcp, saddr = %u, daddr = %u", clskey.saddr, clskey.daddr);
    bpf_printk("tcp, source = %d, dest = %d", bpf_ntohs(clskey.sport) , bpf_ntohs(clskey.dport));
    bpf_printk("tcp, seq = %u, hdrlen = %d, datalen = %d", bpf_ntohl(l4.seq), tcphdrlen, tcp_data_len);
    bpf_printk("*********");

    if (!clsvalue) {
        conn_loss_stat_value clsvalue = {
            .bseq = bpf_ntohl(l4.seq),
            .eseq = bpf_ntohl(l4.seq),
            .elen = tcp_data_len,
            .totalbytes = tcp_data_len,
            .rcvbytes = tcp_data_len,
            .lossbytes = 0,
            .stattime = bpf_ktime_get_ns()
        };
        bpf_map_update_elem(&conn_loss_stats_map, &clskey, &clsvalue, BPF_NOEXIST);
    } else {
        __u32 cseq = bpf_ntohl(l4.seq);
        if (cal_seqoff(clsvalue->eseq, cseq) > clsvalue->elen) {
            bpf_printk("detect loss[cseq = %u, eseq = %u, losslen = %d]", cseq, clsvalue->eseq, cal_seqoff(clsvalue->eseq, cseq) - clsvalue->elen);
            clsvalue->lossbytes += (cal_seqoff(clsvalue->eseq, cseq) - clsvalue->elen);
        }
        if (cal_seqoff(clsvalue->eseq, cseq) >= clsvalue->elen /*|| ((cseq == clsvalue->eseq) && tcp_data_len > 0)*/) {
            clsvalue->eseq = cseq;
            clsvalue->elen = tcp_data_len;
            clsvalue->totalbytes = cal_seqoff(clsvalue->bseq, cseq) + tcp_data_len;
            clsvalue->rcvbytes += tcp_data_len;
            bpf_printk("mmmmmmm = %d --- ggggggggg = %d", clsvalue->totalbytes, clsvalue->rcvbytes);
        }

        bpf_map_update_elem(&conn_loss_stats_map, &clskey, clsvalue, BPF_ANY);

        __u64 ctime = bpf_ktime_get_ns();
        if (ctime - clsvalue->stattime >= 1 * 1000 * 1000 * 1000) {
            // conn_loss_stat_value *clsvalue = bpf_map_lookup_elem(&conn_loss_stats_map, &clskey);
            bpf_printk("9999999999999999999999999");
            bpf_printk("*********");
            bpf_printk("totalbytes = %d, rcvbytes = %d, lossbytes = %d", clsvalue->totalbytes, clsvalue->rcvbytes, clsvalue->lossbytes);
            bpf_printk("*********");
            bpf_map_delete_elem(&conn_loss_stats_map, &clskey);
        }
    }
    return TC_ACT_OK;
}


char g_license[] SEC("license") = "GPL";
